# 参考：https://github.com/Kalafinaian/python-img_gist_feature
import os
from pathlib import Path

import cv2
import sys
import time
import numpy as np

S_NOW_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(S_NOW_DIR)


def img_resize(np_img_in, ln_resize, run_log=None, b_print=False):
    try:
        np_img_resize = cv2.resize(np_img_in, ln_resize, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)
        return np_img_resize, 0
    except Exception as e:
        s_msg = 'resize err:%s' % str(e)
        run_log and run_log.error(s_msg)
        b_print and print(s_msg)
        return None, -3


''' 
Convert raw image to small gray image, resize is  n_resize * n_resize 
'''


def img_2gray(np_img_raw, run_log=None, b_print=False):
    if np_img_raw is None:
        s_msg = "input image null"
        run_log and run_log.error(s_msg)
        b_print and print(s_msg)
        return None, -3
    np_img_gray = None  # Raw Image is BGR imge, so convert rgb to gray
    if len(np_img_raw.shape) == 3 and np_img_raw.shape[2] == 3:
        np_img_gray = cv2.cvtColor(np_img_raw, cv2.COLOR_BGR2GRAY)
    elif len(np_img_raw.shape) == 3 and np_img_raw.shape[2] == 4:  # Raw Image is BGRA imge, there are different situation to solve
        n_sence = 3
        np_img_gray_choose = np.zeros([np_img_raw.shape[0], np_img_raw.shape[1], n_sence], dtype=np.uint8)

        np_img_gray_choose[:, :, 0] = 255 - np_img_raw[:, :, 3]
        np_img_gray_choose[:, :, 1] = cv2.cvtColor(np_img_raw, cv2.COLOR_BGRA2GRAY)
        np_img_gray_choose[:, :, 2] = cv2.cvtColor(np_img_raw[:, :, 0:3], cv2.COLOR_BGR2GRAY)

        ln_sence_non0_num = []  # Get nonzero element of every resize gray
        for i in range(n_sence):
            ln_sence_non0_num.append(len(np_img_gray_choose[:, :, i].nonzero()[0]))
        if len(set(ln_sence_non0_num)) > 1:  # Which image has most nonzero element
            n_max_index = ln_sence_non0_num.index(max(ln_sence_non0_num))
            np_img_gray = np_img_gray_choose[:, :, n_max_index]
        else:  # Which image has most different element
            ln_diff_pix_num = []
            for i in range(n_sence):
                ln_diff_pix_num.append(len(np.unique(np_img_gray_choose[:, :, i])))
            n_max_index = ln_diff_pix_num.index(max(ln_diff_pix_num))
            np_img_gray = np_img_gray_choose[:, :, n_max_index]
    elif len(np_img_raw.shape) == 3 and np_img_raw.shape[2] == 1:  # Raw Image is gray image
        np_img_gray = np_img_raw[:, :, 0]
    elif len(np_img_raw.shape) == 2:
        np_img_gray = np_img_raw
    return np_img_gray, 0


'''  
a image to bgr format 
'''


def img_2bgr(np_img_in, run_log=None, b_print=False):
    if np_img_in is None:
        return None, -3

    np_img_bgr = None  # if raw image is uint16 so conver to uint8
    if len(np_img_in.shape) == 3 and np_img_in.shape[2] == 3:  # Raw Image is BGR imge, so continue
        np_img_bgr = np_img_in
    elif len(np_img_in.shape) == 3 and np_img_in.shape[2] == 4:  # Raw Image is BGRA imge, there are different situation to solve
        h, w, c = np_img_in.shape
        np_img_bgr_1 = cv2.cvtColor(np_img_in, cv2.COLOR_BGRA2BGR)
        b, g, r, a = cv2.split(np_img_in)
        b = cv2.convertScaleAbs(b, alpha=(255.0 / 65535.0))  # (b/256).astype('uint8')
        g = cv2.convertScaleAbs(g, alpha=(255.0 / 65535.0))
        r = cv2.convertScaleAbs(r, alpha=(255.0 / 65535.0))
        a = cv2.convertScaleAbs(a, alpha=(255.0 / 65535.0))
        new_img = cv2.merge((b, g, r))
        not_a = cv2.bitwise_not(a)
        not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR)
        new_img = cv2.bitwise_and(new_img, new_img, mask=a)
        np_img_bgr_2 = cv2.add(new_img, not_a)
        np_img_gray_1 = cv2.cvtColor(cv2.convertScaleAbs(np_img_bgr_1, alpha=(255.0 / 65535.0)), cv2.COLOR_BGR2GRAY)  # Which image has most not white
        np_img_gray_2 = cv2.cvtColor(np_img_bgr_2, cv2.COLOR_BGR2GRAY)
        n_info_1 = len(np.unique(np_img_gray_1))
        n_info_2 = len(np.unique(np_img_gray_2))
        np_img_bgr = np_img_bgr_1 if n_info_1 >= n_info_2 else np_img_bgr_2
    elif len(np_img_in.shape) == 3 and np_img_in.shape[2] == 1:  # Raw Image is gray image
        np_img_bgr = np.tile(np_img_in, (1, 1, 3))  # 256x256x1 ==> 256x256x3
    elif len(np_img_in.shape) == 2:
        np_img_bgr = np.tile(np_img_in, (3, 1, 1))  # 256x256 ==> 3x256x256
        np_img_bgr = np.transpose(np_img_bgr, (1, 2, 0))  # 3x256x256 ==> 256x256x3
    return np_img_bgr, 0


class GistUtils:
    def __init__(self, img_path, n_resize=128, n_w=5, ln_orientation=[8, 8, 8, 8], n_block_num=4, n_prefilt=4):

        self.img = cv2.imread(str(img_path))
        # vector dim(single chanle) = sum(ln_orientation) * n_block_num * n_block_num
        # (8+8+8+8)*(4*4) = 512
        self.n_resize = n_resize
        self.n_boundaryExtension = self.n_resize // 4
        self.n_w = n_w
        self.ln_orientation = ln_orientation
        self.n_block_num = n_block_num  # MUST n_resize % n_block_num == 0
        self.n_prefilt = n_prefilt

        self.__create_gabor()
        self.__get_gfmat()

    def get_gist_vec(self, np_img_raw, mode="rgb"):
        # resize
        np_img_resize, n_ret = img_resize(np_img_raw, (self.n_resize, self.n_resize))
        if n_ret != 0:
            print("image resize error")
            return None

        # convert gray or rgb
        np_gist = None
        if mode.lower() == "gray":
            np_img_gray, n_ret = img_2gray(np_img_resize)
            np_prefilt_img = self.__get_pre_filt(np_img_gray)
            np_gist = self.__gist_main(np_prefilt_img)

        elif mode.lower() == "rgb" or mode.lower() == "bgr":
            np_img_bgr, n_ret = img_2bgr(np_img_resize)

            np_img_b = np_img_bgr[:, :, 0]
            np_img_g = np_img_bgr[:, :, 1]
            np_img_r = np_img_bgr[:, :, 2]

            np_gist_b = self.__get_pre_filt(np_img_b)
            np_gist_g = self.__get_pre_filt(np_img_g)
            np_gist_r = self.__get_pre_filt(np_img_r)

            np_gist_b = self.__gist_main(np_gist_b)
            np_gist_g = self.__gist_main(np_gist_g)
            np_gist_r = self.__gist_main(np_gist_r)

            np_gist = np.hstack([np_gist_b, np_gist_g, np_gist_r])
        else:
            print("input mode error")

        return np_gist.flatten()

    def __get_pre_filt(self, np_img):
        np_log_img = np.log(np_img + 1.0)
        np_pad_img = np.pad(np_log_img, ((self.n_w, self.n_w), (self.n_w, self.n_w)), 'symmetric')

        np_gf = self.np_gf
        np_out = np_pad_img - np.real(np.fft.ifft2(np.fft.fft2(np_pad_img) * np_gf))

        np_local = np.sqrt(np.abs(np.fft.ifft2(np.fft.fft2(np_out ** 2) * np_gf)))
        np_out = np_out / (0.2 + np_local)

        n_size = self.n_resize + 2 * self.n_w

        return np_out[self.n_w: n_size - self.n_w, self.n_w: n_size - self.n_w]

    def __gist_main(self, np_prefilt_img):

        n_b = self.n_boundaryExtension
        np_pad_img = np.pad(np_prefilt_img, ((n_b, n_b), (n_b, n_b)), 'symmetric')
        np_fft2_img = np.fft.fft2(np_pad_img)

        n_filter = self.np_gabor.shape[2]
        n_size = self.np_gabor.shape[0]
        lf_gist = []
        for i in range(n_filter):
            np_res = np.abs(np.fft.ifft2(np_fft2_img * self.np_gabor[:, :, i]))

            np_res = np_res[n_b: n_size - n_b, n_b: n_size - n_b]

            lf_filter_res = self.__down_sampling(np_res)
            lf_gist = lf_gist + lf_filter_res

        np_gist = np.asarray(lf_gist)
        return np_gist[np.newaxis, :]

    def __create_gabor(self):
        n_gabor_size = self.n_resize + 2 * self.n_boundaryExtension
        ln_or = self.ln_orientation

        n_scales = len(ln_or)
        n_filters = sum(ln_or)

        np_param = np.zeros((n_filters, 4), dtype=np.float64)
        n_index = 0
        for i in range(n_scales):
            for j in range(0, ln_or[i]):
                np_param[n_index, 0] = 0.35
                np_param[n_index, 1] = 0.3 / (1.85 ** i)
                np_param[n_index, 2] = 16 * (ln_or[i] ** 2) / (32 ** 2)
                np_param[n_index, 3] = np.pi / ln_or[i] * j

                n_index += 1

        np_linear = np.linspace(-n_gabor_size // 2, n_gabor_size // 2 - 1, n_gabor_size)
        np_fx, np_fy = np.meshgrid(np_linear, np_linear)
        np_res_A = np.fft.fftshift(np.sqrt(np_fx ** 2 + np_fy ** 2))
        np_res_B = np.fft.fftshift(np.angle(np_fx + 1j * np_fy))

        self.np_gabor = np.zeros((n_gabor_size, n_gabor_size, n_filters), dtype=np.float64)
        for i in range(n_filters):
            np_tr = np_res_B + np_param[i, 3]
            np_A = (np_tr < -np.pi) + 0.0
            np_B = (np_tr > np.pi) + 0.0

            np_tr = np_tr + 2 * np.pi * np_A - 2 * np.pi * np_B
            np_every_gabor = np.exp(-10 * np_param[i, 0] * ((np_res_A / n_gabor_size / np_param[i, 1] - 1) ** 2) - 2 * np_param[i, 2] * np.pi * (np_tr ** 2))

            self.np_gabor[:, :, i] = np_every_gabor

    def sav_gist_gabor_to_dir(self, s_gabor_folder='Gabor'):
        if self.np_gabor is None:
            print('gabor is not exists')

        if s_gabor_folder == None or s_gabor_folder == '':
            s_gabor_folder = 'Gabor'

        if not os.path.exists(s_gabor_folder) or not os.path.isdir(s_gabor_folder):
            os.makedirs(s_gabor_folder)

        n_row, n_col, n_num = self.np_gabor.shape
        for i in range(n_num):
            np_submat = self.np_gabor[:, :, i]
            np.savetxt('%s/Gabor%s' % (s_gabor_folder, str(i).zfill(2)), np_submat)
        print('Write Success.')

    def __get_gfmat(self):
        n_s1 = self.n_prefilt / np.sqrt(np.log(2))
        n_boundray = self.n_resize + 2 * self.n_w

        np_linear = np.linspace(-n_boundray // 2, n_boundray // 2 - 1, n_boundray)
        np_fx, np_fy = np.meshgrid(np_linear, np_linear)

        #        np_gf = np.fft.fftshift(np.exp( -(np_fx **2 + np_fy **2)/(n_s1 ** 2)))
        self.np_gf = np.fft.fftshift(np.exp(-(np_fx ** 2 + np_fy ** 2) / (n_s1 ** 2)))

    def __down_sampling(self, np_img):
        np_index = np.linspace(0, self.n_resize, self.n_block_num + 1, dtype=np.int)
        ln_data = []
        for i in range(self.n_block_num):
            for j in range(self.n_block_num):
                np_zone = np_img[np_index[i]: np_index[i + 1], np_index[j]: np_index[j + 1]]
                np_zone = np_zone.T.reshape(-1)
                #                n_res = np.median(np_zone)
                n_res = np.max(np_zone)
                ln_data.append(n_res)
        return ln_data


def image_GIST(img_path):
    """
    给定图片地址，计算GIST特征
    :param img_path:
    :return: 一维向量1536
    """
    gist_helper = GistUtils(img_path)
    np_gist = gist_helper.get_gist_vec(gist_helper.img)
    return np_gist


if __name__ == '__main__':
    img_path = Path(r'E:\dataset\vision\15-Scene\00\3.jpg')
    gist = image_GIST(img_path)
    print(gist.shape)  # 1536维
